package com.sui.bigdata.sml.model.util;

import com.feidee.fd.sml.algorithm.forecast.*;
import com.feidee.fdhadoop.hdfs.HdfsUtils;
import com.sui.bigdata.sml.model.exception.InvalidHdfsPathException;
import com.sui.bigdata.sml.model.model.Model;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.pool2.BaseKeyedPooledObjectFactory;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import scala.collection.JavaConverters;
import scala.collection.Map$;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static com.sui.bigdata.sml.model.util.ModelType.*;

/**
 * @author songhaicheng
 * @date 2019/12/4 14:04
 * @description
 * @reviewer
 */
@Slf4j
public class ModelFactory extends BaseKeyedPooledObjectFactory<Model, Forecast> {

    @Override
    public Forecast create(Model model) throws Exception {
        if (getModelType(model.getModelType()) != TENSORFLOW) {
            if (getModelType(model.getModelType()) == SPARK) {
                for (String path : model.getHdfsPath().split(",")) {
                    checkModel(path);
                }
            } else {
                checkModel(model.getHdfsPath());
            }
        }
        log.info("加载模型：{}", model.toString());
        // 构建 Forecast 对象
        Forecast forecast;
        List<String> pathList = Arrays.asList(model.getHdfsPath().split(Symbol.COMMA.value()));
        // 构建预测字段信息（Java -> Scala）
        scala.collection.immutable.Map<String, String[]> scalaFieldInfoMap = null;
        if (CollectionUtils.isNotEmpty(model.getWarmUpFeature())) {
            Map<String, String[]> fieldInfoMap = new HashMap<>();
            model.getWarmUpFeature().forEach(f -> fieldInfoMap.put(f.getField(), new String[]{f.getDataType(), f.getValue()}));
            scala.collection.mutable.Map<String, String[]> scalaMap = JavaConverters
                    .mapAsScalaMapConverter(fieldInfoMap)
                    .asScala();
            Object obj = Map$.MODULE$.<String, String[]>newBuilder().$plus$plus$eq(scalaMap);
            Object result = ((scala.collection.mutable.Builder) obj).result();
            scalaFieldInfoMap = (scala.collection.immutable.Map) result;
        }

        // 根据模型类型加载 Forecast 对象
        switch (getModelType(model.getModelType())) {
            case SPARK:
                forecast = new SparkForecast(SparkUtil.getLocalSparkSession(), scalaFieldInfoMap,
                        JavaConverters.asScalaIteratorConverter(pathList.iterator()).asScala().toSeq());
                break;
            case PMML:
                forecast = new PMMLForecast(SparkUtil.getLocalSparkSession(), scalaFieldInfoMap,
                        JavaConverters.asScalaIteratorConverter(pathList.iterator()).asScala().toSeq());
                break;
            case MLEAP:
                forecast = new MleapForecast(model.getThreadNum(),
                        scalaFieldInfoMap,
                        JavaConverters.asScalaIteratorConverter(pathList.iterator()).asScala().toSeq());
                break;
            case TENSORFLOW:
                forecast = new TensorflowForecast(Constants.TENSORFLOW_SERVING_IP,
                        Constants.TENSORFLOW_SERVING_PORT,
                        Constants.TENSORFLOW_MODEL_NAME, 1, scalaFieldInfoMap);
                break;
            default:
                throw new Exception("模型类型（modelType）设置错误，错误值为 \"" + model.getModelType() + "\"，请改正");
        }
        return forecast;
    }

    @Override
    public PooledObject<Forecast> wrap(Forecast forecast) {
        return new DefaultPooledObject<>(forecast);
    }

    private void checkModel(String hdfsPath) {
        if (!HdfsUtils.fileExists(hdfsPath)) {
            throw new InvalidHdfsPathException(hdfsPath);
        }
    }

}
